/*
 * @Description: yolo检测头文件
 * @Author: BiChunkai 321521004@qq.com
 * @Date: 2024-10-25 12:08:26
 * @FilePath: /wsMyROS2_TEST/CPP_CODE/TrayDetect/yoloInterface.h
 *
 * Copyright (c) 2024 by 无锡捷普迅科技有限公司, All Rights Reserved.
 */
#ifndef yoloFamily_H
#define yoloFamily_H

#include <string>
#include <iostream>
#include <opencv2/opencv.hpp>
#include <openvino/openvino.hpp> // openvino库
#include <fstream>
#include <vector>
#include <random>

/// @brief 包含模型检测所需的参数
struct Config
{
    float confThreshold;   /// 置信度阈值，决定是否接收检测结果
    float nmsThreshold;    /// 非极大值抑制阈值，减少重叠框
    float scoreThreshold;  /// 用于进一步筛选检测结果的分数阈值
    int inpWidth;          /// 输入图像的宽度
    int inpHeight;         /// 输入图像的高度
    std::string onnx_path; /// ONNX模型文件的路径
};

/// @brief 存储调整大小后的图像和宽高差
struct Resize
{
    cv::Mat resized_image; /// 调整后的图像
    int dw;                /// 宽度差
    int dh;                /// 高度差
};

/// @brief 存储检测结果的信息
struct Detection
{
    int class_id;     /// 检测到的对象类别ID
    float confidence; /// 置信度得分
    cv::Rect box;     /// 矩形框，表示检测到对象的边界
};

/// @brief 输出的检测结果
struct outputDetection
{
    std::string className; /// 检测到的类别名称
    float confidence;      /// 置信度得分
    cv::Rect box;          /// 矩形框
};

/// @brief YOLOv8类
class YOLOV8
{
public:
    /// @brief 构造函数
    /// @param config 模型检测所需要的参数
    YOLOV8(Config config, std::string classPath = "");

    /// @brief 析构
    ~YOLOV8();

    /// @brief 主检测函数，目标检测
    /// @param frame 检测输入图像
    void detect(cv::Mat &frame);

    /// @brief 获取检测参数
    /// @return 检测参数数组
    std::vector<outputDetection> &getDetection();

private:
    // 配置参数
    float confThreshold;
    float nmsThreshold;
    float scoreThreshold;
    int inpWidth;
    int inpHeight;

    // 原始图像与调整图像之间的宽高比
    float rx; // the width ratio of original image and resized image
    float ry; // the height ratio of original image and resized image

    std::string onnx_path;                      /// onnx模型文件名称
    Resize resize;                              /// Resize变量
    ov::Tensor input_tensor;                    /// 输入张量
    ov::InferRequest infer_request;             /// 推理请求
    ov::CompiledModel compiled_model;           /// 编译的模型
    std::vector<outputDetection> outDetections; /// 用于输出的检测参数
    std::string classesPath;                    /// 从txt文件中读取classes

    /// @brief 从文件中加载classname
    /// @param classesPath 检测识别目标的路径
    void loadClassesFromFile(std::string classesPath);

    /// @brief 初始化模型
    /// 1. 读取 ONNX 模型
    /// 2. 配置预处理器
    /// 3. 编译模型并创建推理请求
    void initialmodel();

    /// @brief 图像预处理函数
    ///
    /// 1. 计算输入图像的缩放比例，保持宽高比
    /// 2. 调整图像大小，并在必要时填充图像以适应模型的输入要求
    /// 3. 计算原始图像与调整后图像之间的宽高比，以便后续处理。
    /// 4. 创建一个 OpenVINO 输入张量，并将其设置为推理请求的输入。
    /// 5. 使用异常处理机制确保在出现问题时能够给出适当的错误信息。
    ///
    /// @param frame 输入图像
    void preprocess_img(cv::Mat &frame);

    /// @brief 后处理图像
    ///
    /// 1. 将检测到的物体框绘制到输入图像上
    /// 2. 并显示相应的类别和置信度
    ///
    /// @param frame 输入的图像帧，将在其上绘制检测结果。
    /// @param detections  模型推理得到的检测结果数组。
    /// @param output_shape 输出张量的形状，用于了解检测结果的维度。
    void postprocess_img(cv::Mat &frame, float *detections, ov::Shape &output_shape);

    /// @brief 后处理图像带有返回值
    ///
    /// 1. 将检测到的物体框绘制到输入图像上
    /// 2. 并显示相应的类别和置信度
    ///
    /// @param frame 输入的图像帧，将在其上绘制检测结果。
    /// @param detections 模型推理得到的检测结果数组。
    /// @param output_shape 输出张量的形状，用于了解检测结果的维度。
    void postprocess_img_output(cv::Mat &frame, float *detections, ov::Shape &output_shape);
};

#endif